using Contione.RabbitMQ.Client;
using RabbitMQ.Client;
using RabbitMQ.Client.Events;
using System;
using System.Collections.Concurrent;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace RabbitMQClient
{
    public static class RpcClientExt
    {
        public static RpcClient<T, TResult> StartRpcChannel<T, TResult>(this IRabbitMQClient client)
            where T : class, new()
            where TResult : class, new()
        {
            return new RpcClient<T, TResult>(client);
        }
    }

    public class RpcClient<T, TResult>
            where T : class, new()
            where TResult : class, new()
    {
        private readonly IConnection connection;
        private readonly IModel channel;
        private readonly string replyQueueName;
        private readonly EventingBasicConsumer consumer;

        private readonly ConcurrentDictionary<string, TaskCompletionSource<TResult>> callbackMapper =
                    new ConcurrentDictionary<string, TaskCompletionSource<TResult>>();

        public RpcClient(IRabbitMQClient client)
        {
            connection = client.CreateNewConnection();
            channel = connection.CreateModel();
            replyQueueName = channel.QueueDeclare().QueueName;
            consumer = new EventingBasicConsumer(channel);
            consumer.Received += (model, ea) =>
            {
                if (!callbackMapper.TryRemove(ea.BasicProperties.CorrelationId, out TaskCompletionSource<TResult> tcs))
                {
                    return;
                }
                var body = ea.Body.ToArray();
                var response = Encoding.UTF8.GetString(body);
                TResult result = Newtonsoft.Json.JsonConvert.DeserializeObject<TResult>(response);
                tcs.TrySetResult(result);
            };
        }

        public Task<TResult> CallAsync(string rpc_queue, T request, CancellationToken cancellationToken = default(CancellationToken))
        {
            IBasicProperties props = channel.CreateBasicProperties();
            var correlationId = Guid.NewGuid().ToString();
            props.CorrelationId = correlationId;
            props.ReplyTo = replyQueueName;
            string message = Newtonsoft.Json.JsonConvert.SerializeObject(request);
            var messageBytes = Encoding.UTF8.GetBytes(message);
            var tcs = new TaskCompletionSource<TResult>();
            callbackMapper.TryAdd(correlationId, tcs);

            channel.BasicPublish(
                exchange: "",
                routingKey: rpc_queue,
                basicProperties: props,
                body: messageBytes);

            channel.BasicConsume(
                consumer: consumer,
                queue: replyQueueName,
                autoAck: true);

            cancellationToken.Register(() => callbackMapper.TryRemove(correlationId, out var tmp));
            return tcs.Task;
        }

        public void Close()
        {
            connection.Close();
        }
    }
}